#!/usr/bin/env python
import os
#os.chdir("..")
import torch
import cv2
import argparse
import yaml
import logging 
import numpy as np
import matplotlib.cm as cm
import matplotlib.pyplot as plt
from src.utils.plotting import make_matching_figure

from src.loftr import LoFTR, default_cfg

def make_matching_plot_fast(image0, image1, mkpts0,mkpts1, color,
							margin=10, opencv_title=''):
    H0, W0 = image0.shape
    H1, W1 = image1.shape
    H, W = max(H0, H1), W0 + W1 + margin

    out = 255*np.ones((H, W), np.uint8)
    out[:H0, :W0] = image0
    out[:H1, W0+margin:] = image1
    out = np.stack([out]*3, -1)

    mkpts0, mkpts1 = np.round(mkpts0).astype(int), np.round(mkpts1).astype(int)
    color = (np.array(color[:, :3])*255).astype(int)[:, ::-1]
    for (x0, y0), (x1, y1), c in zip(mkpts0, mkpts1, color):
        c = c.tolist()
        cv2.line(out, (x0, y0), (x1 + margin + W0, y1),
                 color=c, thickness=1, lineType=cv2.LINE_AA)
        # display line end-points as circles
        cv2.circle(out, (x0, y0), 2, c, -1, lineType=cv2.LINE_AA)
        cv2.circle(out, (x1 + margin + W0, y1), 2, c, -1,
                   lineType=cv2.LINE_AA)
    return out
class TrajPlotter(object):
    def __init__(self):
        self.traj = np.zeros((600, 600, 3), dtype=np.uint8)
        pass

    def update(self, est_xyz):
        x, z = est_xyz[0], est_xyz[2]
        est = np.array([x, z]).reshape(2)

        # === drawer ==================================
        # each point
        draw_x, draw_y = int(x) + 290, int(z) + 290
        # draw trajectory
        cv2.circle(self.traj, (draw_x, draw_y), 1, (0, 255, 0), 1)

        return self.traj

# -----------------------------------------------------------------------------
def trans1(image):
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    image = cv2.resize(image, (640, 480))
    return image
def trans2(image):
    image = torch.from_numpy(image)[None][None].cuda() / 255.
    return image
# -----------------------------------------------------------------------------
param_fx = 666.0
param_cx = 665.0
param_cy = 281.0
# -----------------------------------------------------------------------------
absolute_scale = 10.0
traj_plotter = TrajPlotter()
# -----------------------------------------------------------------------------
def run():
    # create camera
    cam = cv2.VideoCapture(0)
    # create matcher
    matcher = LoFTR(config=default_cfg)
    matcher.load_state_dict(torch.load("weights/indoor_ds.ckpt")['state_dict'])
    matcher = matcher.eval().cuda()

    traj_plotter = TrajPlotter()

    last_img = np.zeros([640,480],np.uint8)
    cur_img = np.zeros([640,480],np.uint8)
    cur_R = np.identity(3)
    cur_t = np.zeros((3, 1))

    success, img = cam.read()
    frame_index = 0
    while success:
        success, img = cam.read()
        if frame_index == 0:
            last_img = trans1(img)
        if frame_index > 0:
            cur_img = trans1(img)
            img0 = trans2(last_img)
            img1 = trans2(cur_img)
            batch = {'image0': img0, 'image1': img1}

            # Inference with LoFTR and get prediction
            with torch.no_grad():
                matcher(batch)
                mkpts0 = batch['mkpts0_f'].cpu().numpy()
                mkpts1 = batch['mkpts1_f'].cpu().numpy()
                mconf = batch['mconf'].cpu().numpy()
            # Draw
            color = cm.jet(mconf)
            fig = make_matching_plot_fast(last_img, cur_img, mkpts0, mkpts1, color)
            cv2.imshow('matches', fig)
            E, mask = cv2.findEssentialMat(mkpts1, mkpts0,
                           focal=param_fx, pp=(param_cx, param_cy),
                           method=cv2.RANSAC, prob=0.999, threshold=1.0)
            _, R, t, mask = cv2.recoverPose(E, mkpts1, mkpts0,
                            focal=param_fx, pp=(param_cx, param_cy))
            if (absolute_scale > 0.1):
                cur_t = cur_t + absolute_scale * cur_R.dot(t)
                cur_R = R.dot(cur_R)
        last_img = trans1(img)
        frame_index = frame_index + 1
        # === drawer ==================================
        canvas = traj_plotter.update(cur_t)
        cv2.imshow("trajectory", canvas)
        if cv2.waitKey(1) == 27:
            break


# -----------------------------------------------------------------------------
if __name__ == "__main__":

    run()







